{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W390nLBmZlId"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/learn/search/semantic-search/semantic-search.ipynb) [![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/nbviewer-shield.svg)](https://nbviewer.org/github/pinecone-io/examples/blob/master/learn/search/semantic-search/semantic-search.ipynb)\n",
        "\n",
        "# Semantic Search\n",
        "\n",
        "[![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/fast-link.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/docs/semantic-search.ipynb)\n",
        "\n",
        "In this walkthrough we will see how to use Pinecone for semantic search. To begin we must install the required prerequisite libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q03L1BYEZQfe"
      },
      "outputs": [],
      "source": [
        "!pip install -qU \\\n",
        "  \"pinecone-client[grpc]\"==5.0.1 \\\n",
        "  datasets==2.14.6 \\\n",
        "  sentence-transformers==2.2.2 \\\n",
        "  pinecone-plugin-inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MVSBodoujoZ6"
      },
      "source": [
        "---\n",
        "\n",
        "🚨 _Note: the above `pip install` is formatted for Jupyter notebooks. If running elsewhere you may need to drop the `!`._\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hrSfFiIC5roI"
      },
      "source": [
        "## Data Preprocessing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kujS_e8s55oJ"
      },
      "source": [
        "The dataset preparation process requires a few steps:\n",
        "\n",
        "1. We download the Quora dataset from Hugging Face Datasets.\n",
        "\n",
        "2. The text content of the dataset is embedded into vectors.\n",
        "\n",
        "3. We reformat into a `(id, vector, metadata)` structure to be added to Pinecone.\n",
        "\n",
        "We will see how steps `1`, `2`, and `3` are done in this section, but we won't implement `2` and `3` across the whole dataset until we reach the *upsert loop* as we will iteratively perform these two steps.\n",
        "\n",
        "In either case, this can take some time. If you'd rather skip the data preparation step and get straight to upserts and testing the semantic search functionality, you should\n",
        "refer to the [**fast notebook**](https://colab.research.google.com/github/pinecone-io/examples/blob/master/docs/semantic-search.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IeJPWu9P7EtR"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset('quora', split='train[240000:290000]')\n",
        "dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ngFHND1nQQU2"
      },
      "source": [
        "The dataset contains ~400K pairs of natural language questions from Quora."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CsA67WpW7El4"
      },
      "outputs": [],
      "source": [
        "dataset[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_Zy8zoeQmRZ"
      },
      "source": [
        "Whether or not the questions are duplicates is not so important, all we need for this example is the text itself. We can extract them all into a single `questions` list."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "heGUpy_37Eis"
      },
      "outputs": [],
      "source": [
        "questions = []\n",
        "\n",
        "for record in dataset['questions']:\n",
        "    questions.extend(record['text'])\n",
        "\n",
        "# remove duplicates\n",
        "questions = list(set(questions))\n",
        "print('\\n'.join(questions[:5]))\n",
        "print(len(questions))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BknpfucRkkm"
      },
      "source": [
        "With our questions ready to go we can move on to demoing steps **2** and **3** above.\n",
        "\n",
        "### Building Embeddings and Upsert Format\n",
        "\n",
        "To create our embeddings we will us the `multilingual-e5-large` model from Pinecone inference. We initialize it like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uxcGjb9GSEqA"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pinecone import Pinecone\n",
        "\n",
        "api_key = os.environ.get('PINECONE_API_KEY') or 'PINECONE_API_KEY'\n",
        "pc = Pinecone(api_key=api_key)\n",
        "\n",
        "model_name = \"multilingual-e5-large\"\n",
        "\n",
        "query = 'which city is the most populated in the world?'\n",
        "\n",
        "xq = pc.inference.embed(\n",
        "    model=model_name,\n",
        "    inputs=[query],\n",
        "    parameters={\n",
        "        \"input_type\": \"query\",\n",
        "        \"truncate\": \"END\"\n",
        "    }\n",
        ").data[0]['values']\n",
        "xq[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iy2itPb0S5js"
      },
      "source": [
        "There are *three* interesting bits of information in the above model printout. Those are:\n",
        "\n",
        "* `max_seq_length` is `256`. That means that the maximum number of tokens (like words) that can be encoded into a single vector embedding is `256`. Anything beyond this *must* be truncated.\n",
        "\n",
        "* `word_embedding_dimension` is `1024`. This number is the dimensionality of vectors output by this model. It is important that we know this number later when initializing our Pinecone vector index.\n",
        "\n",
        "* `Normalize()`. This final normalization step indicates that all vectors produced by the model are normalized. That means that models that we would typical measure similarity for using *cosine similarity* can also make use of the *dotproduct* similarity metric. In fact, with normalized vectors *cosine* and *dotproduct* are equivalent.\n",
        "\n",
        "Moving on, we can create a sentence embedding using this model like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z355VLWFqumq"
      },
      "outputs": [],
      "source": [
        "query = 'which city is the most populated in the world?'\n",
        "\n",
        "# create the query vector\n",
        "xq = pc.inference.embed(\n",
        "    model=model_name,\n",
        "    inputs=[query],\n",
        "    parameters={\n",
        "        \"input_type\": \"query\",\n",
        "        \"truncate\": \"END\"\n",
        "    }\n",
        ").data[0]['values']\n",
        "\n",
        "# Print the first 5 vectors as a sanity check\n",
        "xq[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qVZi8xevUWM6"
      },
      "source": [
        "Encoding this single sentence leaves us with a `1024` dimensional sentence embedding (aligned to the `word_embedding_dimension` above).\n",
        "\n",
        "To prepare this for `upsert` to Pinecone, all we do is this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T38HdqxwVg6p"
      },
      "outputs": [],
      "source": [
        "_id = '0'\n",
        "metadata = {'text': query}\n",
        "\n",
        "vectors = [(_id, xq, metadata)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wiXig_rHV2Wz"
      },
      "source": [
        "Later when we do upsert our data to Pinecone, we will be doing so in batches. Meaning `vectors` will be a list of `(id, embedding, metadata)` tuples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ebd7XSamfMsC"
      },
      "source": [
        "## Creating an Index\n",
        "\n",
        "Now the data is ready, we can set up our index to store it.\n",
        "\n",
        "We begin by initializing our connection to Pinecone. To do this we need a [free API key](https://app.pinecone.io)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mc66NEBAcQHY"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# initialize connection to pinecone (get API key at app.pinecone.io)\n",
        "api_key = os.environ.get('PINECONE_API_KEY') or 'PINECONE_API_KEY'\n",
        "\n",
        "# configure client\n",
        "pc = Pinecone(api_key=api_key)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwELn6nBjoaE"
      },
      "source": [
        "Now we setup our index specification, this allows us to define the cloud provider and region where we want to deploy our index. You can find a list of all [available providers and regions here](https://docs.pinecone.io/docs/projects)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l39IjK9xjoaE"
      },
      "outputs": [],
      "source": [
        "from pinecone import ServerlessSpec\n",
        "\n",
        "cloud = os.environ.get('PINECONE_CLOUD') or 'aws'\n",
        "region = os.environ.get('PINECONE_REGION') or 'us-east-1'\n",
        "\n",
        "spec = ServerlessSpec(cloud=cloud, region=region)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SdaTip6CfllN"
      },
      "source": [
        "Now we create a new index called `semantic-search`. It's important that we align the index `dimension` and `metric` parameters with those required by the `MiniLM-L6` model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nzhBBX2PjoaF",
        "tags": [
          "parameters"
        ]
      },
      "outputs": [],
      "source": [
        "index_name = 'semantic-search'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vucJh8sgjoaF"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "# check if index already exists (it shouldn't if this is first time)\n",
        "if index_name not in pc.list_indexes().names():\n",
        "    # if does not exist, create index\n",
        "    pc.create_index(\n",
        "        index_name,\n",
        "        dimension=1024,\n",
        "        metric='cosine',\n",
        "        spec=spec\n",
        "    )\n",
        "    # wait for index to be initialized\n",
        "    while not pc.describe_index(index_name).status['ready']:\n",
        "        time.sleep(1)\n",
        "\n",
        "# connect to index\n",
        "index = pc.Index(index_name)\n",
        "# view index stats\n",
        "index.describe_index_stats()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YUd1VGg6i108"
      },
      "source": [
        "Now we upsert the data, we will do this in batches of `128`.\n",
        "\n",
        "_**Note:** On Google Colab with GPU expected runtime is ~7 minutes. If using CPU this will be significantly longer. If you'd like to get this running faster refer to the [fast notebook](https://github.com/pinecone-io/examples/blob/master/search/semantic-search/semantic-search-fast.ipynb)._"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7o6rdgU8olqd"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "batch_size = 96\n",
        "vector_limit = 100000\n",
        "\n",
        "questions = questions[:vector_limit]\n",
        "\n",
        "for i in tqdm(range(0, len(questions), batch_size)):\n",
        "    # find end of batch\n",
        "    i_end = min(i+batch_size, len(questions))\n",
        "    # create IDs batch\n",
        "    ids = [str(x) for x in range(i, i_end)]\n",
        "    # create metadata batch\n",
        "    metadatas = [{'text': text} for text in questions[i:i_end]]\n",
        "    # create embeddings\n",
        "    xc = pc.inference.embed(\n",
        "        model=model_name,\n",
        "        inputs=questions[i:i_end],\n",
        "        parameters={\n",
        "            \"input_type\": \"passage\",\n",
        "            \"truncate\": \"END\"\n",
        "        }\n",
        "    ).data\n",
        "    embeddings = [item['values'] for item in xc]\n",
        "    # create records list for upsert\n",
        "    records = zip(ids, embeddings, metadatas)\n",
        "    # upsert to Pinecone\n",
        "    index.upsert(vectors=records)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VrK_IN079Vuu"
      },
      "source": [
        "## Making Queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rr4unPAq9alb"
      },
      "source": [
        "Now that our index is populated we can begin making queries. We are performing a semantic search for *similar questions*, so we should embed and search with another question. Let's begin."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SNa4WHoJ0COt"
      },
      "outputs": [],
      "source": [
        "query = \"which city has the largest population in the world?\"\n",
        "\n",
        "# create the query vector\n",
        "xq = pc.inference.embed(\n",
        "    model=model_name,\n",
        "    inputs=[query],\n",
        "    parameters={\n",
        "        \"input_type\": \"query\",\n",
        "        \"truncate\": \"END\"\n",
        "    }\n",
        ").data[0]['values']\n",
        "\n",
        "# now query\n",
        "xc = index.query(vector=xq, top_k=5, include_metadata=True)\n",
        "xc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-XwOWcgo_QtI"
      },
      "source": [
        "In the returned response `xc` we can see the most relevant questions to our particular query. We can reformat this response to be a little easier to read:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gy7isg_f-vWg"
      },
      "outputs": [],
      "source": [
        "for result in xc['matches']:\n",
        "    print(f\"{round(result['score'], 2)}: {result['metadata']['text']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1JK5yApl_5fE"
      },
      "source": [
        "These are good results, let's try and modify the words being used to see if we still surface similar results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dJbjE-iq_yMr"
      },
      "outputs": [],
      "source": [
        "query = \"which metropolis has the highest number of people?\"\n",
        "\n",
        "# create the query vector\n",
        "xq = pc.inference.embed(\n",
        "    model=model_name,\n",
        "    inputs=[query],\n",
        "    parameters={\n",
        "        \"input_type\": \"query\",\n",
        "        \"truncate\": \"END\"\n",
        "    }\n",
        ").data[0]['values']\n",
        "\n",
        "# now query\n",
        "xc = index.query(vector=xq, top_k=5, include_metadata=True)\n",
        "for result in xc['matches']:\n",
        "    print(f\"{round(result['score'], 2)}: {result['metadata']['text']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HIAxOPb-A2w_"
      },
      "source": [
        "Here we used different terms in our query than that of the returned documents. We substituted **\"city\"** for **\"metropolis\"** and **\"populated\"** for **\"number of people\"**.\n",
        "\n",
        "Despite these very different terms and *lack* of term overlap between query and returned documents — we get highly relevant results — this is the power of *semantic search*.\n",
        "\n",
        "You can go ahead and ask more questions above. When you're done, delete the index to save resources:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-cWdeKzhAtww"
      },
      "outputs": [],
      "source": [
        "pc.delete_index(index_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PMMJSu_DbRx0"
      },
      "source": [
        "---"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
